#!/usr/bin/env python
import numpy as np

class LinearRegression():
    """Linear regression model.

    y = X @ w
    t ~ N(t|X @ w, var)
    """

    def fit(self, x_train, y_train):
        """Perform least squares fitting.

        Parameters
        ----------
        x_train: (N, D) np.ndarray
            training independent variable 
        y_train: (N,  ) np.ndarray
            training dependent variable 
        """
        # A = x_train.T @ x_train

        # self.w = np.linalg.solve(A + 1e-7 * np.identity(A.shape[0]), x_train.T @ y_train)
        # self.w = np.linalg.solve(A, x_train.T @ y_train)
        self.w, _, _, _ = np.linalg.lstsq(x_train, y_train, rcond=None, )
        self.var = np.mean(np.square(x_train @ self.w - y_train))

    def predict(self, x):
        """Return prediction given input.

        Parameters
        ----------
        x : (N, D) np.ndarray
            samples to predict their output 
        return_std : bool, optional
            returns standard deviation of each predition if True

        Returns
        -------
        y : np.ndarray
            prediction of each sample (N,) 
        """
        y = x @ self.w
        return y

    
